import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import numpy as np
import copy

# 构建Embedding类来实现文本嵌入层
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        # d_model:词嵌入的维度
        # vocab:词表的大小
        super(Embeddings, self).__init__()
        # 定义Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        # 将参数传入类中
        self.d_model = d_model

    def forward(self, x):
        # x:代表输入进模型的文本通过词汇映射后的数字张量
        return self.lut(x) * math.sqrt(self.d_model)

# 构建位置编码器的类
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        # d_model:代表词嵌入的维度
        # dropout:代表dropout层的置零比率
        # max_len：代表每个句子的最大长度
        super(PositionalEncoding, self).__init__()

        # 实例化dropout层
        self.dropout = nn.Dropout(p=dropout)

        # 初始化一个位置编码矩阵，大小是max_len*d_model
        pe = torch.zeros(max_len, d_model)

        # 初始化一个绝对位置矩阵，max_len * 1
        position = torch.arange(0, max_len).unsqueeze(1)

        # 定义一个变化矩阵div_term，跳跃式的初始化
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))

        # 将前面定义的变化矩阵进行奇数，偶数的分别赋值
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数列分别赋值
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数列分别赋值

        # 将二维张量扩充成三维张量
        pe = pe.unsqueeze(0)  # 从第0维进行扩充

        # 将位置编码矩阵注册成模型的buffer,这个buffer不是模型中的参数，不跟随优化器同步更新
        # 注册成buffer后我们就可以在模型保存后重新加载的时候，将这个位置编码器和模型参数一同加载进来
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x:代表文本序列的词嵌入表示
        # 首先明确pe的编码太长了，将第二个维度，也就是max_len对应的那个维度缩小成x的句子长度
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)


# 构建掩码张量的函数
def subsequent_mask(size):
    # size：代表掩码张量后两个维度，形成一个方阵
    attn_shape = (1, size, size)

    # 使用np.ones()先构建一个全1的张量，然后利用np.triu()形成上三角矩阵
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')

    # 使得这个三角矩阵反转
    return torch.from_numpy(1 - subsequent_mask)

def attention(query, key, value, mask=None, dropout=None):
    # query，key,value代表注意力的三个输入张量
    # mask:掩码张量，
    # dropout是nn.Dropout层的实例化对象，默认为None

    # 在函数中，首先取query的最后一维的大小，一般情况下就等同于我们的词嵌入维度，命名为d_k
    d_k = query.size(-1)
    # 按照注意力公式，将query与key的转置相乘，这里面key是将最后两个维度进行转置，再除以缩放系数根号下d_k进行缩放
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    # 接着判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法，将掩码张量和scores张量每个位置一一比较，如果掩码张量处为0，则使用-1e9（比较小的数据代替）
        # 则对应的scores张量用-1e9这个值来替换，如果下演示
        scores = scores.masked_fill(mask == 0, -1e9)

    # 对scores的最后一维进行softmax操作，使用F.softmax方法，每一个参数是softmax对象，第二个是目标维度
    # 这样获得最终的注意力张量
    p_attn = F.softmax(scores, dim=-1)

    # 之后判断是否使用dropout进行随机置0
    if dropout is not None:
        # 将p_attn传入dropout对象进行'丢弃'处理
        p_attn = dropout(p_attn)

    # 最后，根据公式将p_attn与value张量相乘获得最终的query注意力表示，同时返回注意力张量
    return torch.matmul(p_attn, value), p_attn

# 实现克隆函数，因为在多头注意力机制下，要用到多个结构相同的线性层
# 需要使用clone函数将它们一同初始化到一个网络层列表对象中
def clone(module, N):
    # module:代表要克隆目标网络层
    # N：将module克隆几个
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


# 实现多头注意力机构的类
class MultiHeadedAttention(nn.Module):
    def __init__(self, head, embedding_dim, dropout=0.1):
        # header:代表几个头的参数
        # embedding_dim:代表词嵌入的维度
        # dropout:进行dropout操作时，置零的比率
        super(MultiHeadedAttention, self).__init__()

        # 要确认一个事实：多头的数量header需要整除词入的维度embedding_dim
        assert embedding_dim % head == 0
        # 得到每个头获得词向量的维度
        self.d_k = embedding_dim // head
        self.head = head
        self.embedding_dim = embedding_dim

        # 获得线性层，要获得4个，分别是Q,K,V以及最终输出线性层
        self.linears = clone(nn.Linear(embedding_dim, embedding_dim), 4)

        # 初始化注意力张量
        self.attn = None

        # 初始化dropout对象
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        # query，key,value代表注意力的三个输入张量
        # mask:掩码张量，
        # 首先判断是否使用掩码张量
        if mask is not None:
            # 使用squeeze将掩码张量进行维度扩充，代表多头中的第n个头
            mask = mask.unsqueeze(0)
        # 得到batch_size
        batch_size = query.size(0)

        # 首先使用zip将网络层和输入数据连接在一起，模型的输出利用view和transpose进行维度和形状的改变
        query, key, value = \
            [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
             for model, x in zip(self.linears, (query, key, value))]

        # 将每个头的输出传入到注意力层
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)

        # 得到每个头的计算结果是4维张量,需要进行形状的转换
        # 前面已经将1，2两个维度进行过转置，在这里要重新转置回来
        # 注意：经历了transpose()方法后，必须要使用contiguous方法，不然无法使用view()方法
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)
        # 最后将x输入线性层列表中的最后一个线性层中进行处理，得到最终的多头注意力结构输出
        return self.linears[-1](x)


# 通过类PositionwiseFeedForward来实现前馈全连接层
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        # 初始化函数有三个输入参数分别是d_model，d_ff,dropout，第一个是线性层的输入维度也是第二个线性层的输出维度，
        # 因为我们希望输入通过前馈全连接层后输入和输出的维度不变，第二个参数d_ff就是第二个线性层的输入维度和第一个线性层的输出
        # 最后一个是dropout置0比率
        super(PositionwiseFeedForward, self).__init__()

        # 首先按照我们预期使用nn实例化了两个线性层对象，self.w1和self.w2
        # 它们的参数分别是d_model,d_ff和d_ff,d_model
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        # 然后使用nn的Dropout实例化了对象self.dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 首先经过第一个线性层，然后使用Functional中relu函数进行激活，
        # 之后再使用dropout进行随机置0，最后通过第二个性性层w2,返回最终结果
        return self.w2(self.dropout(F.relu(self.w1(x))))

# 通过layerNorm实现规范化层的类
class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        # 初始化函数有两个参数，一个是features,表示词嵌入的维度，
        # 另一个是eps它是一个足够小的数，在规范化公式的分母中出现，
        # 防止分母为0
        super(LayerNorm, self).__init__()

        # 根据features的形状初始化两个参数张量a2和b2，第一个初始化为1张量，
        # 也就是里面的元素是1，第二个初始化为0张量，也就是里面的元素都是0，这两个张量就是规范化的参数
        # 因为直接对上一层得到的结果做规范化公式计算，将改变结果的正常表征，因此就需要有参数作为调节因子
        # 使其既能满足规范化要求，又能不改变针对目标的表征，最后使用nn.parameter封装，代表他们是模型的参数。
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))

        # 把eps传到类中
        self.eps = eps

    def forward(self, x):
        # 输入参数x代表来自上一层的输出
        # 在函数中，首先对输入变量x求其最后一个维度的均值，并保持输出维度与输入维度一致
        # 接着再求最后一个维度的标准差，然后就是根据规范化公式，用x减去均值除以标准差获得规范化的结果
        # 最后对结果乘以我们的缩放参数，即a2,*号代表同型点乘，即对应位置进行乘法操作，加上位移参数b2，返回即可
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

# 使用SublayerConnection来实现子层连接结构的类
class SublayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        # 它输入参数有两个，size以及dropout,size一般都是词嵌入维度的大小，
        # dropout本身是对模型结构中的节点数进行随机抑制的比率，
        # 又因为节点限被抑制等效就是该节点的输出都是0，因此也可以把dropout看作是对输出矩阵的随机值
        super(SublayerConnection, self).__init__()
        # 实例化了规范化对象self.norm
        self.norm = LayerNorm(size)
        # 又使用nn中预定义的dropout实例化一个self.dropout对象。
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, sublayer):
        # 前向逻辑函数中，接收上一个层或者子层的输入作为第一个参数
        # 将该子层连接中的子层函数作为第二个参数
        # 我们首先对输出进行规范化，然后将结果传给子层处理，之后再对子层进行dropout操作，
        # 随机停止一些网络中神经元的作用，来防止过拟合，最后还有一个add操作
        # 因为存在跳跃连接，所以是将输入x与dropout后的子层输出结果相加作为最终的子层连接输出
        return x + self.dropout(sublayer(self.norm(x)))

# 使用EncoderLayer类实现编码层
class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        # size是词嵌入维度大小，self_atn是多头自注意力子层，feed_forward是前馈全连接层
        super(EncoderLayer, self).__init__()

        # 首先将self_attn和feed_forward传入其中
        self.self_attn = self_attn
        self.feed_forward = feed_forward

        # 编码器层中有两个子层连接结构，所以使用clones函数进行克隆
        self.sublayer = clone(SublayerConnection(size, dropout), 2)

        self.size = size

    def forward(self, x, mask):
        # x是上一层的输出，mask是掩码张量
        # 第一个子层包含多头自注意力子层，第二个子层包含前馈全连接子层，最后返回结果
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

# 使用Encoder类来实现编码器
class Encoder(nn.Module):
    def __init__(self, layer, N):
        #layer是指编码器层，N是指编码器层的个数
        super(Encoder, self).__init__()
        # 首先使用clone函数克隆N个编码器层放在self.layers中
        self.layers = clone(layer, N)
        # 再初始化一个规范化层，它将用在编码器的最后面
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        # x代表上一层的输出，mask代表掩码张量
        # 首先对克隆的编码器层进行循环，每次都会得到一个新的x,
        # 这个循环的过程，就相当于输出的x经过了N个编码器层的处理
        # 最后再通过规范化层的对象self.norm进行处理，最后返回结果
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

# 使用DecoderLayer的类实现解码器层
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        # size代表词嵌入的维度大小，同时也代表解码器层的尺寸
        # self_attn，多头注意力对象，也就是说这个注意力机制需要Q=K=V，
        # src_attn，多头注意力对象，这里的Q!=K=V
        # feed_forward，前鐀全连接层对象
        super(DecoderLayer, self).__init__()

        # 在初始化函数中，主要就是将这些输入传到类中
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        # 按照结构图使用clone函数克隆三个子层连接对象
        self.sublayer = clone(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, source_mask, target_mask):
        # x是来自上一层的输入x，
        # memory是来自上一层的编码器层的语义存储变量，
        # source_mask和target_mask分别是源数据掩码张量和目标数据掩码张量
        m = memory

        # 将x传入第一个子层结构，第一个子层结构的输入分别是x和self_attn函数，因为是自注意力机制，所以Q,K,V都是x
        # 最后一个参数是目标数据掩码张量，这时要对目标数据进行遮掩，因为此时模型可能还没有生成任何目标数据，
        # 比如在解码器准备生成第一个字符或词汇时，我们其实已经传入了第一个字符以便计算损失，
        # 但是我们不希望在生成第一个字符时模型能利用这个信息，因此我们会将其遮掩，同样生成第二个字符或词汇时,
        # 模型只能使用第一个字符或词汇信息，第二个字符以及之后的信息都不允许被模型使用。
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))

        # 接着进入第二个子层，这个子层中常规的注意力机制，q是输入x;k,v是编码层输出memery,
        # 同样也传入source_mask,但是进行源数据遮掩的原因并非是抑制信息泄漏，而是遮蔽掉对比结果没有意义的字符而产生的注意力值，
        # 以此提升模型效果和训练速度，这样就完成了第二个子层的处理
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))

        # 最后一个子层就是前鐀全连接子层，经过它的处理后就可以返回结果，这就是我们的解码器层结构
        return self.sublayer[2](x, self.feed_forward)

# 构建解码器类
class Decoder(nn.Module):
    def __init__(self, layer, N):
        # layer:代表解码器层的对象
        # N:代表将layer进行几层的拷贝
        super(Decoder, self).__init__()

        # 利用clone函数克隆N个layer
        self.layers = clone(layer, N)
        # 实例化一个规范化层
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, source_mask, target_mask):
        # x是来自上一层的输入x，
        # memory是来自上一层的编码器层的语义存储变量，
        # source_mask和target_mask分别是源数据掩码张量和目标数据掩码张量
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)

# 将线性层和softmax计算层一起实现，因为二者的共同目标是生成最后的结构
# 因此把类的名字叫做Generator，生成器类
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        # 初始化函数的输入参数有两个，d_model代表词嵌入维度，vocab_size代表词表大小
        super(Generator, self).__init__()
        # 首先就是使用nn中的预定义线性层进行实例化，得到一个对象self.project等待使用，
        # 这个线性层的参数有两个，就是初始化函数传进来的两个参数：d_model,vocab_size
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        # 前向逻辑函数中输入是上一层的输出张量x
        # 在函数中，首先使用上一步得到的self.project对x进行线性变化，
        # 然后使用F中已经实现的log_softmax进行softmax处理，
        # 在这里之所以使用log_softmax是因为和我们这个pytorch版本的损失函数实现有关，在其它版本
        # log_softmax就是对softmax的结果又取了对数，因为对数函数是单调递增函数，
        # 因此对最终我们取最大的概率值没有影响，最后返回结果即可
        return F.log_softmax(self.project(x), dim=-1)

# 使用EncoderDecoder类来实现编码-解码器结构
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        # encoder是编码器对象,
        # decoder是解码器对象，
        # source_embed是源数据嵌入函数，
        # target_embed是目标数据嵌入函数
        super(EncoderDecoder, self).__init__()
        # 将参数传入到类中
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator

    def forward(self, source, target, source_mask, target_mask):
        # source代表源数据
        # target代表目标数据
        # source_mask和target_mask代表对应的掩码张量
        # 在函数中，将source,source_mask传入编码函数，得到结果后，
        # 与source_mask,target,和target_mask一同传给解码函数
        return self.generator(
            self.decode(self.encode(source, source_mask), source_mask, target, target_mask)
        )

    def encode(self, source, source_mask):
        # 使用src_embed对source做处理，然后和source_mask一起传给self.encoder
        return self.encoder(self.src_embed(source), source_mask)

    def decode(self, memory, source_mask, target, target_mask):
        # 使用tgt_embed对target做处理，然后和source_mask,target_mask,memory一起传给self.decoder
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)